import pandas as pd
from statsmodels.tsa.seasonal import seasonal_decompose
import matplotlib.pyplot as plt

# ==================== 配置参数 ====================
INPUT_FILE = "ori_message.xlsx"  # 原始数据文件路径（可替换为绝对路径）
OUTPUT_FILE = "monthly_average_data.xlsx"  # 输出文件路径


# ==================== 数据处理函数（保留原逻辑） ====================
def process_data():
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']
    all_monthly_data = []
    for year in years:
        df = excel_file.parse(year)
        df['年'] = df['年'].ffill().astype(int)
        df['月'] = df['月'].ffill().astype(int)
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'
        }).reset_index()
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })
        all_monthly_data.append(grouped)
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})
    result.to_excel(OUTPUT_FILE, index=False)
    return result


# ==================== 季节分解与绘图函数（残差子图改为散点图） ====================
def plot_seasonal_decompose(data):
    # 构造日期列（确保格式正确）
    data['日期'] = data['年'].astype(str) + '-' + data['月'].astype(str) + '-01'
    data['日期'] = pd.to_datetime(data['日期'])
    data = data.set_index('日期')
    ts = data['流量(m3/s)']

    # 季节分解
    decomposition = seasonal_decompose(ts, model='multiplicative', period=12)
    trend = decomposition.trend
    seasonal = decomposition.seasonal
    residual = decomposition.resid  # 残差数据（时间序列）

    # 创建子图（4行1列）
    plt.figure(figsize=(10, 12))

    # 原始数据（折线图）
    plt.subplot(4, 1, 1)
    plt.plot(ts)
    plt.title('Multiplicative Decompose')
    plt.ylabel('流量(m³/s)')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-01'))

    # 趋势（折线图）
    plt.subplot(4, 1, 2)
    plt.plot(trend)
    plt.ylabel('趋势')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-01'))

    # 季节（折线图）
    plt.subplot(4, 1, 3)
    plt.plot(seasonal)
    plt.ylabel('季节')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-01'))

    # 残差（散点图，关键修改）
    plt.subplot(4, 1, 4)
    plt.scatter(  # 替换为scatter函数
        x=residual.index,  # 横轴为时间索引
        y=residual.values,  # 纵轴为残差值
        s=10,  # 散点大小（可调整）
        alpha=0.6,  # 透明度（避免重叠遮挡）
        color='red'  # 颜色（可选）
    )
    plt.ylabel('残差')
    plt.xlabel('年份')
    plt.xlim(pd.Timestamp('2017-01-01'), pd.Timestamp('2021-12-01'))

    plt.tight_layout()
    plt.savefig('流量时间序列分解图.png', dpi=300)
    plt.show()


# ==================== 主程序 ====================
if __name__ == "__main__":
    processed_data = process_data()
    plot_seasonal_decompose(processed_data)
    print("\n========== 每月平均流量、含沙量及排沙量 ==========\n")
    print(processed_data.to_string(index=False))